import os
os.environ['DISPLAY'] = ':99'
os.environ['SDL_VIDEODRIVER'] = 'dummy'  # Add this line to handle authorization error
import gym
from nes_py.wrappers import JoypadSpace
import gym_super_mario_bros
from gym_super_mario_bros.actions import SIMPLE_MOVEMENT,COMPLEX_MOVEMENT
import matplotlib.pyplot as plt
from skimage.transform import resize    
import numpy as np

def downscale_obs(obs,new_size=(42,42),method='bilinear',to_gray=False):
    if to_gray:
        return resize(obs,new_size,anti_aliasing=True).max(axis=2)
    
    return resize(obs,new_size,anti_aliasing=True)


env = gym_super_mario_bros.make('SuperMarioBros-v0')
env = JoypadSpace(env, SIMPLE_MOVEMENT)
done = True
frames = []
for step in range(3000):
    if done:
        state = env.reset()
    state,reward,done,info = env.step(env.action_space.sample())
    frame = env.render(mode="rgb_array")
    frames.append(frame)
    
    
env.close()

img = frames[200]

plt.imshow(downscale_obs(img))
plt.savefig('icm.png')

